#include "concatd_custom_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"
#include <algorithm>

namespace optiling {
const uint32_t BLOCK_DIM = 8;
const uint32_t TILE_NUM = 8;
const uint32_t BUFFER_NUM = 2;
const uint32_t BLOCK_SIZE = 32;        // 基础对齐块大小（字节）
static inline uint64_t AlignUp(uint64_t x, uint64_t a){ return (x + a - 1) / a * a; }
static ge::graphStatus TilingFunc(gert::TilingContext *context)
{
    ConcatdCustomTilingData tiling;
    uint64_t ubLength = 0;
    uint64_t bigCoreDataNum = 0;
    uint64_t bigCoreLoopNum = 0;
    uint64_t bigCoreTailDataNum = 0;

    auto attrs = context->GetAttrs();
    int32_t d = 0;  // 默认值
    if (attrs) {
        const int64_t* attrA = attrs->GetInt(0);  
        if (attrA != nullptr) {
            d = *attrA;
        }
    }
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    const auto xShape = context->GetInputTensor(0)->GetOriginShape();
    const auto yShape = context->GetInputTensor(1)->GetOriginShape();
    const auto inputDataType = context->GetInputTensor(0)->GetDataType();
    context->SetBlockDim(BLOCK_DIM);
    uint32_t typeSize = sizeof(float);
    switch (inputDataType) {
        case ge::DT_FLOAT16: 
            typeSize = sizeof(uint16_t); 
            break;
        case ge::DT_FLOAT:   
            typeSize = sizeof(float);    
            break;
        case ge::DT_INT32:   
            typeSize = sizeof(int32_t);  
            break;
        case ge::DT_INT16:   
            typeSize = sizeof(int16_t);  
            break;
        case ge::DT_UINT8:   
            typeSize = sizeof(uint8_t);
            break;
        default:
                // 不支持的数据类型
            return ge::GRAPH_FAILED;
        }
    uint32_t totalLengthx = context->GetInputShape(0)->GetOriginShape().GetShapeSize();
    uint32_t totalLengthy = context->GetInputShape(1)->GetOriginShape().GetShapeSize();
    uint32_t totalLengthz = totalLengthy + totalLengthx;
    int32_t dimNum = xShape.GetDimNum();
    uint32_t x1 = xShape.GetDim(dimNum - 2);  // x的行数    
    uint32_t x2 = xShape.GetDim(dimNum - 1);   //x的列数    
    uint32_t y1= yShape.GetDim(dimNum - 2);  // y的行数  
    uint32_t y2= yShape.GetDim(dimNum - 1);  // y的列数  
    printf("%d",x2);
    printf("%d",y2);
    uint64_t ubSize = 192 * 1024;           // 使用固定的192KB UB大小
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);

    //核间划分//1
    uint32_t big_core_num = x1 % BLOCK_DIM; 
    uint32_t small_core_num=BLOCK_DIM - big_core_num;
    uint32_t small_tile_length = x1/BLOCK_DIM; //小核分到的行数
    uint32_t big_tile_length = x1/BLOCK_DIM + 1; //大核分到的行数
    //0
    uint32_t sbig_core_num = (x1+y1) % BLOCK_DIM; 
    uint32_t ssmall_core_num=BLOCK_DIM - big_core_num;
    uint32_t ssmall_tile_length = (x1+y1)/BLOCK_DIM; //小核分到的行数
    uint32_t sbig_tile_length = (x1+y1)/BLOCK_DIM + 1; //大核分到的行数

    //核内划分//1维度的张量
    int64_t core_tile_x1 = 1; // 对行维度进行划分
    auto FitsUB = [&](int64_t x2,int64_t y2,int64_t core_tile_x1) -> bool {
        uint64_t xBytes  =  AlignUp((uint64_t)core_tile_x1  * x2 * typeSize, BLOCK_SIZE);
        uint64_t yBytes  =  AlignUp((uint64_t)y2 * core_tile_x1 * typeSize, BLOCK_SIZE);
        // 总 UB = 双缓冲 * (x + z + y + temp)
        uint64_t total = BUFFER_NUM * (xBytes + yBytes) * 2;
         // 预留 5% 余量
        return total <= (ubSize * 95 / 100);
    };

        // 调整核内分块参数以适应UB
        while (FitsUB(x2,y2,core_tile_x1) && core_tile_x1 < big_tile_length) {
            core_tile_x1 *= 2;
        }
        if (core_tile_x1 != 1) { // 判断是否倍增了 回溯到倍增前的结果
            core_tile_x1 /= 2;
        }
        if (!FitsUB(x2,y2,core_tile_x1)) {
            return ge::GRAPH_FAILED;
        }//
    int64_t core_tile_s1 = 1; // 对行维度进行划分

            // 调整核内分块参数以适应UB
        while (FitsUB(x2,y2,core_tile_s1) && core_tile_s1 < sbig_tile_length) {
            core_tile_s1 *= 2;
        }
        if (core_tile_s1 != 1) { // 判断是否倍增了 回溯到倍增前的结果
            core_tile_s1 /= 2;
        }
        if (!FitsUB(x2,y2,core_tile_s1)) {
            return ge::GRAPH_FAILED;
        }//

    uint32_t small_tile_times = small_tile_length/core_tile_x1; 
    uint32_t big_tile_times = big_tile_length/core_tile_x1; 

    uint32_t small_tail_num= small_tile_length % core_tile_x1;
    uint32_t big_tail_num= big_tile_length % core_tile_x1;
    //0
    uint32_t ssmall_tile_times = ssmall_tile_length/core_tile_s1; 
    uint32_t sbig_tile_times = sbig_tile_length/core_tile_s1; 

    uint32_t ssmall_tail_num= ssmall_tile_length % core_tile_s1;
    uint32_t sbig_tail_num= sbig_tile_length % core_tile_s1;

    //1
    if(small_tail_num!=0){
        small_tile_times ++;
    }else{
        small_tail_num = core_tile_x1;
    }
    if(big_tail_num!=0){
        big_tile_times ++;
    }else{
        big_tail_num = core_tile_x1;
    }//0
    if(ssmall_tail_num!=0){
        ssmall_tile_times ++;
    }else{
        ssmall_tail_num = core_tile_s1;
    }
    if(sbig_tail_num!=0){
        sbig_tile_times ++;
    }else{
        sbig_tail_num = core_tile_s1;
    }
    //0维度划分
    //每次搬相同行数，x和y搬的次数不一样
    //uint32_t tile_times = std::max(small_tile_times,big_tile_times);
    //tiling.set_tile_times(tile_times);
    tiling.set_small_tile_times(small_tile_times);
    tiling.set_big_tile_times(big_tile_times);
    tiling.set_small_tail_num(small_tail_num);
    tiling.set_big_tail_num(big_tail_num);
    tiling.set_totalLengthx (totalLengthx);
    tiling.set_totalLengthy(totalLengthy);
    tiling.set_totalLengthz(totalLengthz);
    tiling.set_x1(x1);
    tiling.set_x2(x2);
    tiling.set_y1(y1);
    tiling.set_y2(y2);
    tiling.set_big_core_num(big_core_num);
    tiling.set_small_core_num(small_core_num);
    tiling.set_small_tile_length(small_tile_length);
    tiling.set_big_tile_length(big_tile_length);
    tiling.set_core_tile_x1(core_tile_x1);
    tiling.set_tileNum(TILE_NUM);
//0
    tiling.set_ssmall_tile_times(ssmall_tile_times);
    tiling.set_sbig_tile_times(sbig_tile_times);
    tiling.set_ssmall_tail_num(ssmall_tail_num);
    tiling.set_sbig_tail_num(sbig_tail_num);
    tiling.set_sbig_core_num(sbig_core_num);
    tiling.set_ssmall_core_num(ssmall_core_num);
    tiling.set_ssmall_tile_length(ssmall_tile_length);
    tiling.set_sbig_tile_length(sbig_tile_length);
    tiling.set_core_tile_s1(core_tile_s1);

    tiling.set_d(d);



    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
}
} // namespace optiling



namespace ge {
static graphStatus InferShape(gert::InferShapeContext *context)
{
    // const gert::Shape *x1_shape = context->GetInputShape(0);
    // gert::Shape *y_shape = context->GetOutputShape(0);
    // *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}

static graphStatus InferDataType(gert::InferDataTypeContext *context)
{
    const auto inputDataType = context->GetInputDataType(0);
    context->SetOutputDataType(0, inputDataType);
    return ge::GRAPH_SUCCESS;
}
} // namespace ge

namespace ops {
class ConcatdCustom : public OpDef {
public:
    explicit ConcatdCustom(const char *name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16,ge::DT_FLOAT,ge::DT_INT32,ge::DT_INT16,ge::DT_INT8,ge::DT_INT64,ge::DT_BOOL})
            .Format({ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND});
        this->Input("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16,ge::DT_FLOAT,ge::DT_INT32,ge::DT_INT16,ge::DT_INT8,ge::DT_INT64,ge::DT_BOOL})
            .Format({ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND});
        this->Output("z")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16,ge::DT_FLOAT,ge::DT_INT32,ge::DT_INT16,ge::DT_INT8,ge::DT_INT64,ge::DT_BOOL})
            .Format({ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND});

        this->Attr("d")
            .AttrType(OPTIONAL)
            .Int();

        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);
        
        this->AICore()
            .SetTiling(optiling::TilingFunc)
            .AddConfig("ascend910")
            .AddConfig("ascend310p")
            .AddConfig("ascend310b")
            .AddConfig("ascend910b");
    }
};
OP_ADD(ConcatdCustom);
} // namespace ops
